module DirectApproach where

import Courcoubetis
import Data.Array
import Data.List

----------------------------------------------------
----------------------------------------------------
-- Data type for HMMs
data HMM a = HMM {
		statesHMM	:: Array Int a,			-- States
		pTransHMM	:: Array (Int, Int) Float,	-- State Transition Matrix
		labelFHMM	:: Array Int [String],  	-- Labeling Function
		obsHMM		:: [Char],			-- Set of observation
		pObsHMM		:: Array (Int, Int) Float,	-- Observaton Probability Matrix
		initDiHMM	:: Array Int Float		-- Initial Distribution Matrix
		} deriving (Show)

-- Values of a sample HMM
{-estadosHMM :: Array Int Int
estadosHMM = array (1,2) [(1,1) , (2,2)]

matrixHMM :: Array (Int, Int) Float
matrixHMM = array ((1,1), (2,2)) 
				[((1,1), 0.7), ((1,2), 0.3),
				 ((2,1), 0.4), ((2,2), 0.6)]

etiqHMM :: Array Int [String]
etiqHMM = array (1,2) [(1,["p"]), (2,["q"])]

myObsHMM :: [Char]
myObsHMM = "SML"

obsMatHMM :: Array (Int,Int) Float
obsMatHMM = array ((1,1),(2,3))
				[((1,1), 0.1), ((1,2), 0.4), ((1,3), 0.5),
				 ((2,1), 0.7), ((2,2), 0.2), ((2,3), 0.1)]

initHMM :: Array Int Float
initHMM = array (1,2) [(1, 0.6), (2, 0.4)]
-}
{-estadosHMM :: Array Int Int
estadosHMM = array (1,5) [(1,1) , (2,2), (3,3), (4,4), (5,5)]

matrixHMM :: Array (Int, Int) Float
matrixHMM = array ((1,1), (5,5)) 
				[((1,1), 0), ((1,2), 0.5), ((1,3),   0), ((1,4),   0), ((1,5), 0.5),
				 ((2,1), 0), ((2,2),   0), ((2,3), 0.5), ((2,4),   0), ((2,5), 0.5),
                                 ((3,1), 0), ((3,2),   0), ((3,3),   0), ((3,4), 0.5), ((3,5), 0.5),
                                 ((4,1), 0), ((4,2),   0), ((4,3),   0), ((4,4), 0.5), ((4,5), 0.5),
                                 ((5,1), 0), ((5,2),   0), ((5,3),   0), ((5,4),   0), ((5,5),   1)
                                 ]

etiqHMM :: Array Int [String]
etiqHMM = array (1,5) [(1,["a"]), (2,["a"]), (3,[""]), (4,["b"]), (5,[""])]

myObsHMM :: [Char]
myObsHMM = "SML"

obsMatHMM :: Array (Int,Int) Float
obsMatHMM = array ((1,1),(5,3))
				[((1,1), 1/3), ((1,2), 1/3), ((1,3), 1/3),
				 ((2,1), 1/3), ((2,2), 1/3), ((2,3), 1/3),
                                 ((3,1), 1/3), ((3,2), 1/3), ((3,3), 1/3),
                                 ((4,1), 1/3), ((4,2), 1/3), ((4,3), 1/3),
                                 ((5,1), 1/3), ((5,2), 1/3), ((5,3), 1/3)]

initHMM :: Array Int Float
initHMM = array (1,5) [(1, 1), (2, 1), (3, 1), (4, 1), (5, 1)]
-}
estadosHMM :: Array Int Int
estadosHMM = array (1,4) [(1,1) , (2,2), (3,3), (4,4)]

matrixHMM :: Array (Int, Int) Float
matrixHMM = array ((1,1), (4,4)) 
				[((1,1), 0), ((1,2),   1), ((1,3),   0), ((1,4),   0),
				 ((2,1), 0), ((2,2),0.01), ((2,3),0.01), ((2,4),0.98),
                                 ((3,1), 1), ((3,2),   0), ((3,3),   0), ((3,4),   0),
                                 ((4,1), 0), ((4,2),   0), ((4,3),   0), ((4,4),   1)
                                 ]

etiqHMM :: Array Int [String]
etiqHMM = array (1,4) [(1,[""]), (2,["try"]), (3,["fail"]), (4,["succ"])]

myObsHMM :: [Char]
myObsHMM = "U"

obsMatHMM :: Array (Int,Int) Float
obsMatHMM = array ((1,1),(4,1))
				[((1,1), 1),
				 ((2,1), 1),
                                 ((3,1), 1),
                                 ((4,1), 1)]

initHMM :: Array Int Float
initHMM = array (1,4) [(1, 1), (2, 1), (3, 1), (4, 1)]

myHMM = HMM {statesHMM = estadosHMM, pTransHMM = matrixHMM, labelFHMM = etiqHMM, 
				obsHMM = myObsHMM, pObsHMM = obsMatHMM, initDiHMM = initHMM}

----------------------------------------------------
----------------------------------------------------
-- Data type for POCTL*
-- State Formulas
-- To remove the grammar's ambiguity I do not have a data type for Belief States Formulas, instead I
-- add its expressive power to the POCTL* data type that appears next
data POCTL  = VerdadP | FalsoP | AtomP String | NoP POCTL | OP POCTL POCTL | YP POCTL POCTL | Prob String Float PathF
instance Show POCTL where
	show = poctlS

poctlS :: POCTL -> String
poctlS VerdadP		= "Verdad"
poctlS FalsoP 		= "Falso"
poctlS (AtomP s)	=  s
poctlS (NoP p)		= "~" ++ poctlS p
poctlS (OP p1 p2)	= "(" ++ (poctlS p1) ++ " v " ++ (poctlS p2) ++ ")"
poctlS (YP p1 p2)       = "(" ++ (poctlS p1) ++ " ^ " ++ (poctlS p2) ++ ")"
poctlS (Prob c f path)	= "Pr [" ++ c ++ show f ++ "] (" ++ show path ++ ")"

-- Path Formulas
data PathF  = FormP POCTL | NoPath PathF | OPath PathF PathF | YPath PathF PathF
		| NextPath String PathF | UntilBPath Int PathF PathF | UntilPath PathF PathF

instance Show PathF where
	show = pathS

pathS :: PathF -> String

pathS (FormP phi)	  = poctlS phi
pathS (NoPath p)	  = "~(" ++ pathS p ++ ")"
pathS (OPath p1 p2)	  = "(" ++ (pathS p1) ++ " v " ++ (pathS p2) ++ ")"
pathS (YPath p1 p2)	  = "(" ++ (pathS p1) ++ " ^ " ++ (pathS p2) ++ ")"
pathS (NextPath obs p)	  = "X_{" ++ obs ++ "} (" ++ pathS p ++ ")"
pathS (UntilBPath n p1 p2)= "(" ++ (pathS p1) ++ " U<=" ++ show n ++ " " ++ (pathS p2) ++ ")"
pathS (UntilPath p1 p2)   = "(" ++ (pathS p1) ++ " U " ++ (pathS p2) ++ ")"

formula :: POCTL
--formula = YP (NoP (AtomP "b")) (Prob "<" 0.05 (UntilPath (FormP (AtomP "a")) (NextPath "S" (FormP (AtomP "b")))) )
--formula = YP (NoP (AtomP "b")) (Prob "<" 0.05 (UntilBPath 10000 (FormP (AtomP "a")) (FormP (AtomP "b"))))
formula = Prob ">" 0.98 (UntilBPath 2 (FormP VerdadP) (FormP (AtomP "succ")))

----------------------------------------------------
----------------------------------------------------
-- This function transforms an HMM into a DTMC.
-- Here we perform the transformation from an HMM to a DTMC, so the new 
-- state set is S'=S x O, P'((s,o),(s',o')) = P(s,s').Mu(s',o') and L'(s,o)=L(s)u{Omega c O|o in Omega} 
-- Notice two things: First, the labeling function works as previously has done it, the observation part
-- is going to be added later on. Second, the initial distribution function has been modified according 
-- to the new states
satQOSHMM :: HMM a -> DTMC (Courcou (a,Char))
satQOSHMM hmm = DTMC {statesDTMC = getStates, 			
                      pTransDTMC = newTrans, 			
                      labelFDTMC = newLabel, 					
                      initDiDTMC = initialT}		
  where	
    tmpList	= [Base (s,o) | s<-elems (statesHMM hmm), o <- (obsHMM hmm)]	
    tmpLength	= length tmpList	
    obsLen	= length (obsHMM hmm)	
    getStates	= listArray (1, tmpLength) tmpList	
    newTrans	= array ((1,1), (tmpLength, tmpLength)) 
                  [((n,m), (pTransHMM hmm)!((div (n-1) obsLen)+1, (div (m-1) obsLen)+1) * 
                           (pObsHMM hmm)!((div (m-1) obsLen)+1, (mod (m-1) obsLen)+1))				
                  | n<-[1..tmpLength], m<-[1..tmpLength]]	
    newLabel 	= array (1, tmpLength) [(i, (labelFHMM hmm)!((div (i-1) obsLen)+1))                                     
                                       | i<-[1..tmpLength]]	
    initialT 	= array (1, tmpLength) [(i, (initDiHMM hmm)!((div (i-1) obsLen)+1) * 
                                            (pObsHMM hmm)!((div (i-1) obsLen)+1, (mod (i-1) obsLen)+1)) 	
                                       | i<-[1..tmpLength]]

----------------------------------------------------
----------------------------------------------------
-- Function that transforms QOS formula into QLS fromulas
transForm :: PathF -> LTL
transForm (FormP (AtomP s))	= Atom s
transForm (FormP VerdadP)	= Verdad
transForm (FormP FalsoP)	= Falso
transForm (NoPath p) 		= No (transForm p)
transForm (OPath p q)		= O (transForm p) (transForm q)
transForm (YPath p q)           = Y (transForm p) (transForm q)
transForm (NextPath s p)	= Y (Atom s) (Next (transForm p))
transForm (UntilBPath n p q)    = UntilB n (transForm p) (transForm q)
transForm (UntilPath p q)       = Until (transForm p) (transForm q)
transForm _			= error "This case shouldn't happen. Because the formula we are considering is a maximal state subformula"

----------------------------------------------------
----------------------------------------------------
-- This function returns the very first component of the state
-- produced by transforming the HMM to a DTMC
origin :: Courcou (a, Char) -> a
origin (Base (n, _)) 	= n
origin (Par (cou, _))	= origin cou

----------------------------------------------------
----------------------------------------------------
-- Auxiliary function that determines if the satisfying probability
-- is within the reqired range

comparison :: Float -> String -> Float -> Bool
comparison x str y = case str of
			"<=" -> x <= y
			"<"  -> x < y
			">=" -> x >= y
			otherwise -> x > y -- The only case remaining is str = ">"

----------------------------------------------------
----------------------------------------------------
-- This method computes the satisfying probability of HMM state 'i' 
-- given the DTMC m and the set of DTMC states that satisfy the specification,
-- this corresponds to a summation of initial distributions of m
probVal :: Eq a => DTMC (Courcou (a, Char)) -> a -> [Courcou (a, Char)] -> Float
probVal m i sat = sum [(initDiDTMC m)!j | j<-indices dtmcStat, elem (dtmcStat!j) sat, 
						origin(dtmcStat!j) == i]
							where
								dtmcStat	= statesDTMC m

----------------------------------------------------
----------------------------------------------------
-- This function computes the HMM states whose satisfying probability of the LTL formula
-- is within the range establish by 'c' and 'x', the comparison operator and the real number
-- between 0 and 1, for this end we use the DTMC resulting of the HMM transformation
satisfy :: Eq a => HMM a -> DTMC (Courcou (a, Char)) -> LTL -> String -> Float -> [a]
satisfy hmm _ Verdad c x        = case c of 
								">"  | x<1  -> elems (statesHMM hmm)
								">="        -> elems (statesHMM hmm)
								"<=" | x==1 -> elems (statesHMM hmm)
								otherwise -> []
satisfy hmm _ Falso  c x        = case c of 
								"<"  | x>0  -> elems (statesHMM hmm) 
								"<="        -> elems (statesHMM hmm)
								">=" | x==0 -> elems (statesHMM hmm)
								otherwise -> []
satisfy hmm m (Atom s) c x      = [hmmStat!i | i <- indices hmmStat, 
								comparison (probVal m (hmmStat!i) (sat m (Atom s))) c x]
									where
										hmmStat 	= statesHMM hmm
satisfy hmm m (No p) c x        = [hmmStat!i | i <- indices hmmStat, 
								comparison (probVal myMC (hmmStat!i) (sat myMC ap)) c x]
									where 
										hmmStat 	= statesHMM hmm
										(myMC, _) 	= cxu m (No p) 0
										(ap, _)		= reWrite (No p) 0
satisfy hmm m (O p1 p2) c x     = [hmmStat!i | i <- indices hmmStat, 
								comparison (probVal myMC (hmmStat!i) mySet) c x]
									where 
										hmmStat 	   = statesHMM hmm
										mySet 		   = union (sat myMC reWp1) (sat myMC reWp2)
										(myMC, _)	   = cxu m (O p1 p2) 0
										(O reWp1 reWp2, _) = reWrite (O p1 p2) 0
satisfy hmm m (Y p1 p2) c x     = [hmmStat!i | i <- indices hmmStat, 
								comparison (probVal myMC (hmmStat!i) mySet) c x]
									where 
										hmmStat            = statesHMM hmm
										mySet 	           = intersect (sat myMC reWp1) (sat myMC reWp2)
										(myMC, _)	   = cxu m (Y p1 p2) 0
										(Y reWp1 reWp2, _) = reWrite (Y p1 p2) 0
satisfy hmm m (Next p) c x	= [hmmStat!i | i <- indices hmmStat, 
								comparison (probVal myMC (hmmStat!i) (sat myMC ap)) c x]	
									where
										hmmStat 	= statesHMM hmm
										(myMC, _)	= cxu m (Next p) 0
										(ap, _)		= reWrite (Next p) 0
satisfy hmm m (UntilB n p1 p2) c x = [hmmStat!i | i <- indices hmmStat, 
								comparison (probVal myMC (hmmStat!i) (sat myMC ap)) c x]	
									where
										hmmStat 	= statesHMM hmm
										(myMC, _)	= cxu m (UntilB n p1 p2) 0
										(ap, _)		= reWrite (UntilB n p1 p2) 0
satisfy hmm m (Until p1 p2) c x = [hmmStat!i | i <- indices hmmStat, 
								comparison (probVal myMC (hmmStat!i) (sat myMC ap)) c x]	
									where
										hmmStat 	= statesHMM hmm
										(myMC, _)	= cxu m (Until p1 p2) 0
										(ap, _)		= reWrite (Until p1 p2) 0
                                                                                                  
----------------------------------------------------
----------------------------------------------------
-- It returns the set of HMM's states satisfying the POCTL* path formula obtained from 
-- the probabilistic operator
directApp :: Eq a => HMM a -> PathF -> String -> Float -> [a]
directApp hmm p str x = satisfy hmm (satQOSHMM hmm) (transForm p) str x